#include "des.h"

uint64_t initial_permutation(uint64_t input) {
    uint64_t result = 0;
    for (int i = 0; i < 64; i++) {
        result <<= 1;
        result |= (input >> (64 - IP[i])) & LB64_MASK;
    }
    return result;
}

uint64_t inverse_initial_permutation(uint64_t input) {
    uint64_t result = 0;
    for (int i = 0; i < 64; i++) {
        result <<= 1;
        result |= (input >> (64 - PI[i])) & LB64_MASK;
    }
    return result;
}

void generate_subkeys(uint64_t key, uint64_t sub_key[16]) {
    uint64_t permuted_choice_1 = 0;
    for (int i = 0; i < 56; i++) {
        permuted_choice_1 <<= 1;
        permuted_choice_1 |= (key >> (64 - PC1[i])) & LB64_MASK;
    }

    uint32_t C = (permuted_choice_1 >> 28) & 0x0FFFFFFF;
    uint32_t D = permuted_choice_1 & 0x0FFFFFFF;

    for (int i = 0; i < 16; i++) {
        for (int j = 0; j < iteration_shift[i]; j++) {
            C = ((C << 1) | (C >> 27)) & 0x0FFFFFFF;
            D = ((D << 1) | (D >> 27)) & 0x0FFFFFFF;
        }

        uint64_t CD = ((uint64_t)C << 28) | D;
        sub_key[i] = 0;
        for (int j = 0; j < 48; j++) {
            sub_key[i] <<= 1;
            sub_key[i] |= (CD >> (56 - PC2[j])) & LB64_MASK;
        }
    }
}

uint32_t feistel(uint32_t R, uint64_t subkey) {
    uint64_t expanded = 0;
    for (int j = 0; j < 48; j++) {
        expanded <<= 1;
        expanded |= (R >> (32 - E[j])) & LB32_MASK;
    }

    expanded ^= subkey;

    uint32_t s_output = 0;
    for (int j = 0; j < 8; j++) {
        uint64_t mask = (expanded >> (42 - 6*j)) & 0x3F;
        char row = ((mask & 0x20) >> 4) | (mask & 0x01);
        char col = (mask >> 1) & 0x0F;

        s_output <<= 4;
        s_output |= S[j][16 * row + col] & 0x0F;
    }

    uint32_t result = 0;
    for (int j = 0; j < 32; j++) {
        result <<= 1;
        result |= (s_output >> (32 - P[j])) & LB32_MASK;
    }

    return result;
}

uint64_t des_encrypt_block(uint64_t block, uint64_t subkeys[16]) {
    uint64_t ip = initial_permutation(block);
    uint32_t L = (uint32_t)(ip >> 32) & L64_MASK;
    uint32_t R = (uint32_t)(ip & L64_MASK);

    for (int i = 0; i < 16; i++) {
        uint32_t temp = R;
        R = L ^ feistel(R, subkeys[i]);
        L = temp;
    }

    uint64_t pre_output = ((uint64_t)R << 32) | L;
    return inverse_initial_permutation(pre_output);
}

uint64_t des_decrypt_block(uint64_t block, uint64_t subkeys[16]) {
    uint64_t ip = initial_permutation(block);
    uint32_t L = (uint32_t)(ip >> 32) & L64_MASK;
    uint32_t R = (uint32_t)(ip & L64_MASK);

    for (int i = 0; i < 16; i++) {
        uint32_t temp = R;
        R = L ^ feistel(R, subkeys[15 - i]);
        L = temp;
    }

    uint64_t pre_output = ((uint64_t)R << 32) | L;
    return inverse_initial_permutation(pre_output);
}


/*
 * The DES function
 * input: 64 bit message
 * key: 64 bit key for encryption/decryption
 * mode: 'e' = encryption; 'd' = decryption
 */
uint64_t des(uint64_t input, uint64_t key, char mode) {
    uint64_t subkeys[16] = {0};
    generate_subkeys(key, subkeys);

    if (mode == 'd') {
        return des_decrypt_block(input, subkeys);
    } else {
        return des_encrypt_block(input, subkeys);
    }
}
